#include "hnswlib.h"
#include "HNSW_Louvain.h"
#include <iostream>
#include <cstdio>
#include <cstdlib>
#include <vector>
#include <unordered_set>
#include <cstring>
#include <sys/stat.h>
#include <chrono>

using namespace std;

float *fvecs_read(const char *fname, size_t *d_out, size_t *n_out) {
    FILE *f = fopen(fname, "r");
    if (!f) {
        fprintf(stderr, "could not open %s\n", fname);
        perror("");
        abort();
    }
    int d;
    fread(&d, 1, sizeof(int), f);
    fseek(f, 0, SEEK_SET);
    struct stat st;
    fstat(fileno(f), &st);
    size_t sz = st.st_size;
    size_t n = sz / ((d + 1) * 4);

    *d_out = d;
    *n_out = n;
    float *x = new float[n * (d + 1)];
    size_t nr = fread(x, sizeof(float), n * (d + 1), f);


    // shift array to remove row headers
    for (size_t i = 0; i < n; i++)
        memmove(x + i * d, x + 1 + i * (d + 1), d * sizeof(*x));


    fclose(f);
    return x;
}

// not very clean, but works as long as sizeof(int) == sizeof(float)
int *ivecs_read(const char *fname, size_t *d_out, size_t *n_out) {
    return (int *) fvecs_read(fname, d_out, n_out);
}

class TimeRecord {
    std::chrono::steady_clock::time_point time_begin;
public:
    TimeRecord() {
        time_begin = std::chrono::steady_clock::now();
    }

    float get_elapsed_time_micro() {
        std::chrono::steady_clock::time_point time_end = std::chrono::steady_clock::now();
        return (std::chrono::duration_cast<std::chrono::microseconds>(time_end - time_begin).count());
    }

    void reset() {
        time_begin = std::chrono::steady_clock::now();
    }

};

double count_recall(int n_query, int topk, std::vector<std::unordered_set<int>> test_gnd_l, int *gnd,int dim) {

    std::vector<double> recall_l;
    double avg_recall = 0;

    for (int i = 0; i < n_query; i++) {
        int match = 0;
        for (int j = 0; j < topk; j++) {
            if (test_gnd_l[i].find(gnd[i * dim + j]) != test_gnd_l[i].end()) {
                match++;
            }
        }
        double recall = (double) match / topk;
        recall_l.push_back(recall);
        avg_recall += recall;
    }
    return avg_recall / n_query;
}


using namespace hnswlib;

int main() {
    vector<vector<vector<pair<int, vector<int>>>>> partitions;
    size_t d, nq, nd;
    size_t M = 16;
    size_t efConstruction = 100;
    size_t dim;
    float *query = fvecs_read("D:\\ClionProject\\HNSW\\siftsmall_query.fvecs", &d, &nq);
    float *data = fvecs_read("D:\\ClionProject\\HNSW\\siftsmall_base.fvecs", &d, &nd);
    int *answer = ivecs_read("D:\\ClionProject\\HNSW\\siftsmall_groundtruth.ivecs", &dim, &nq);
    auto *space = new hnswlib::L2Space(d);
    vector<long> degree;
    vector<int> links;
    vector<float> weight;
    long m = 0;
    {
        auto *hnsw = new hnswlib::HierarchicalNSW<float>(space, nd, M, efConstruction, 100);
        int report_every = 5000;
        TimeRecord insert_record;

        for (int i = 0; i < nd; i++) {
            if (i % report_every == 0) {
                std::cout << i / (0.01 * nd) << " %, "
                          << 1e-6 * insert_record.get_elapsed_time_micro() << " s/iter" << " \n";
                insert_record.reset();
            }
            hnsw->addPoint(data + i * d, i);
        }
        //vector<float> avg_dis;
        for (int i = 0; i < nd; i++) {
            unsigned int *neighbor = hnsw->get_linklist0(i);
            long deg = *neighbor;
            m += deg;
            degree.push_back(m);
            // float totalDis=0;
            for (int j = 1; j <= deg; j++) {
                int neighID = *(neighbor + j);
                links.push_back(neighID);
                float nowDis = 0;
                for (int t = 0; t < d; t++) {
                    nowDis += pow(*(data + i * d + t) - *(data + neighID * d + t), 2);
                }
                weight.push_back(1 / log(nowDis));
                // totalDis+=nowDis;
            }
            // totalDis=totalDis/(float)deg;
            //avg_dis.push_back(totalDis);
        }
        printf("Have built HNSW\n");
        delete hnsw;
    }

    HNSW_Louvain hl(data, d, nd, m, degree, links, weight);
    int topk = 10;
    const int ef_search_config[3] = {10, 5, 15}; //start, stride, iter
    for (int iter = 0, ef_search = ef_search_config[0];
         iter < ef_search_config[2]; iter++, ef_search += ef_search_config[1]) {
        vector<priority_queue < myPair, vector<myPair>, ComparatorForMaxHeap >>
                result;
        float *tmp_query = query;
        TimeRecord query_record;
        for (int i = 0; i < nq; i++, tmp_query += d) {
            priority_queue <myPair, vector<myPair>, ComparatorForMaxHeap> W = hl.knn(tmp_query, topk, ef_search);
            result.push_back(W);
        }
        double query_time = query_record.get_elapsed_time_micro() / nq * 1e-3; //unit ms

        vector<unordered_set<int>> test_gnd_l(nq);
        for (int i = 0; i < nq; i++) {
            priority_queue <myPair, vector<myPair>, ComparatorForMaxHeap> tmp_pq = result[i];
            while (!tmp_pq.empty()) {
                myPair p = tmp_pq.top();
                tmp_pq.pop();
                test_gnd_l[i].insert(p.first);
            }
        }
        double recall = count_recall(nq, topk, test_gnd_l, answer,dim);
        printf("efSearch %d, average recall: %.3f, query time %.3f ms\n", ef_search, recall, query_time);
    }

    return 0;

}
